# -*- coding: utf-8 -*-

import tensorflow as tf


def _seg_net_convolution(in_put, in_channel, out_channel, layer_name, is_training=True):
    with tf.variable_scope(layer_name) as scope:
        weights = tf.get_variable(name="weights", shape=[3, 3, in_channel, out_channel],
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        biases = tf.get_variable(name="biases", shape=[out_channel],
                                 initializer=tf.constant_initializer())
        weight_decay = tf.mul(tf.nn.l2_loss(weights), 0.00001, name="weight_decay_loss")
        tf.add_to_collection("weight_loss", weight_decay)
        convolution = tf.nn.bias_add(value=tf.nn.conv2d(input=in_put,
                                                        filter=weights,
                                                        strides=[1, 1, 1, 1],
                                                        padding="SAME"),
                                     bias=biases)
        bn = tf.contrib.layers.batch_norm(inputs=convolution,
                                          decay=0.9,
                                          is_training=is_training,
                                          updates_collections=None)
        output = tf.nn.relu(bn, name=scope.name)
    return output


def _seg_net_max_pool(in_put):
    values = tf.nn.max_pool(value=in_put,
                            ksize=[1, 2, 2, 1],
                            strides=[1, 2, 2, 1],
                            padding="SAME")
    return values


def _seg_net_up_sample(in_put, in_channel, out_channel, layer_name):
    with tf.variable_scope(layer_name) as scope:
        weights = tf.get_variable(name="weights", shape=[2, 2, in_channel, out_channel],
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        weight_decay = tf.mul(tf.nn.l2_loss(weights), 0.00001, name="weight_decay_loss")
        tf.add_to_collection("weight_loss", weight_decay)
        input_shape = in_put.get_shape().as_list()
        assert len(input_shape) == 4
        output_shape = [input_shape[0], input_shape[1]*2, input_shape[2]*2, input_shape[3]]
        output = tf.nn.conv2d_transpose(value=in_put,
                                        filter=weights,
                                        output_shape=output_shape,
                                        strides=[1, 2, 2, 1],
                                        name=scope.name)
        return output


def seg_net_network(in_put, is_training=True):
    # encoder
    conv1_1 = _seg_net_convolution(in_put=in_put, in_channel=3, out_channel=64,
                                   layer_name="conv1_1", is_training=is_training)
    conv1_2 = _seg_net_convolution(in_put=conv1_1, in_channel=64, out_channel=64,
                                   layer_name="conv1_2", is_training=is_training)
    pool_1_values = _seg_net_max_pool(conv1_2)
    conv2_1 = _seg_net_convolution(in_put=pool_1_values, in_channel=64, out_channel=128,
                                   layer_name="conv2_1", is_training=is_training)
    conv2_2 = _seg_net_convolution(in_put=conv2_1, in_channel=128, out_channel=128,
                                   layer_name="conv2_2", is_training=is_training)
    pool_2_values = _seg_net_max_pool(conv2_2)
    conv3_1 = _seg_net_convolution(in_put=pool_2_values, in_channel=128, out_channel=256,
                                   layer_name="conv3_1", is_training=is_training)
    conv3_2 = _seg_net_convolution(in_put=conv3_1, in_channel=256, out_channel=256,
                                   layer_name="conv3_2", is_training=is_training)
    conv3_3 = _seg_net_convolution(in_put=conv3_2, in_channel=256, out_channel=256,
                                   layer_name="conv3_3", is_training=is_training)
    pool_3_values = _seg_net_max_pool(conv3_3)
    conv4_1 = _seg_net_convolution(in_put=pool_3_values, in_channel=256, out_channel=512,
                                   layer_name="conv4_1", is_training=is_training)
    conv4_2 = _seg_net_convolution(in_put=conv4_1, in_channel=512, out_channel=512,
                                   layer_name="conv4_2", is_training=is_training)
    conv4_3 = _seg_net_convolution(in_put=conv4_2, in_channel=512, out_channel=512,
                                   layer_name="conv4_3", is_training=is_training)
    pool_4_values = _seg_net_max_pool(conv4_3)
    conv5_1 = _seg_net_convolution(in_put=pool_4_values, in_channel=512, out_channel=512,
                                   layer_name="conv5_1", is_training=is_training)
    conv5_2 = _seg_net_convolution(in_put=conv5_1, in_channel=512, out_channel=512,
                                   layer_name="conv5_2", is_training=is_training)
    conv5_3 = _seg_net_convolution(in_put=conv5_2, in_channel=512, out_channel=512,
                                   layer_name="conv5_3", is_training=is_training)
    pool_5_values = _seg_net_max_pool(conv5_3)
    # decoder
    pool5_d = _seg_net_up_sample(pool_5_values, in_channel=512, out_channel=512, layer_name="pool5_d")
    conv5_3_d = _seg_net_convolution(in_put=pool5_d, in_channel=512, out_channel=512,
                                     layer_name="conv5_3_d", is_training=is_training)
    conv5_2_d = _seg_net_convolution(in_put=conv5_3_d, in_channel=512, out_channel=512,
                                     layer_name="conv5_2_d", is_training=is_training)
    conv5_1_d = _seg_net_convolution(in_put=conv5_2_d, in_channel=512, out_channel=512,
                                     layer_name="conv5_1_d", is_training=is_training)
    pool4_d = _seg_net_up_sample(conv5_1_d, in_channel=512, out_channel=512, layer_name="pool4_d")
    conv4_3_d = _seg_net_convolution(in_put=pool4_d, in_channel=512, out_channel=512,
                                     layer_name="conv4_3_d", is_training=is_training)
    conv4_2_d = _seg_net_convolution(in_put=conv4_3_d, in_channel=512, out_channel=512,
                                     layer_name="conv4_2_d", is_training=is_training)
    conv4_1_d = _seg_net_convolution(in_put=conv4_2_d, in_channel=512, out_channel=256,
                                     layer_name="conv4_1_d", is_training=is_training)
    pool3_d = _seg_net_up_sample(conv4_1_d, in_channel=256, out_channel=256, layer_name="pool3_d")
    conv3_3_d = _seg_net_convolution(in_put=pool3_d, in_channel=256, out_channel=256,
                                     layer_name="conv3_3_d", is_training=is_training)
    conv3_2_d = _seg_net_convolution(in_put=conv3_3_d, in_channel=256, out_channel=256,
                                     layer_name="conv3_2_d", is_training=is_training)
    conv3_1_d = _seg_net_convolution(in_put=conv3_2_d, in_channel=256, out_channel=128,
                                     layer_name="conv3_1_d", is_training=is_training)
    pool2_d = _seg_net_up_sample(conv3_1_d, in_channel=128, out_channel=128, layer_name="pool2_d")
    conv2_2_d = _seg_net_convolution(in_put=pool2_d, in_channel=128, out_channel=128,
                                     layer_name="conv2_2_d", is_training=is_training)
    conv2_1_d = _seg_net_convolution(in_put=conv2_2_d, in_channel=128, out_channel=64,
                                     layer_name="conv2_1_d", is_training=is_training)
    pool1_d = _seg_net_up_sample(conv2_1_d, in_channel=64, out_channel=64, layer_name="pool1_d")
    conv1_2_d = _seg_net_convolution(in_put=pool1_d, in_channel=64, out_channel=64,
                                     layer_name="conv1_2_d", is_training=is_training)
    with tf.variable_scope("conv1_1_d"):
        weights = tf.get_variable(name="weights", shape=[3, 3, 64, 2],
                                  initializer=tf.contrib.layers.xavier_initializer_conv2d())
        biases = tf.get_variable(name="biases", shape=[2],
                                 initializer=tf.constant_initializer())
        weight_decay = tf.mul(tf.nn.l2_loss(weights), 0.00001, name="weight_decay_loss")
        tf.add_to_collection("weight_loss", weight_decay)
        conv1_1_d = tf.nn.bias_add(value=tf.nn.conv2d(input=conv1_2_d,
                                                      filter=weights,
                                                      strides=[1, 1, 1, 1],
                                                      padding="SAME"),
                                   bias=biases)
    return conv1_1_d


def main():
    with tf.Graph().as_default():
        input_x = tf.placeholder(dtype=tf.float32, shape=[1, 480, 256, 3])
        result = seg_net_network(input_x)
        print(result.get_shape())


if __name__ == "__main__":
    main()
